/*
 * Description: performance registry controll utils
 */
#include "registry_utils_ex.h"
#include "log.h"

namespace PERF {

#ifdef UNICODE
// Special characters
const TCHAR CHAR_TAIL = L'\0';
const TString STR_EMPTY = L"";
// Registry Error Description
const TString ERR_UNKNOWN = L"unknown error";
#else

// Special characters
const TCHAR CHAR_TAIL = '\0';
const TString STR_EMPTY = "";
// Registry Error Description
const TString ERR_UNKNOWN = "unknown error";
#endif

const uint32_t ERROR_STRING_SIZE = 101;

RegistryUtilsEx::RegistryUtilsEx(HKEY hRoot, const TString& targetKey) noexcept
    : rootKey_(hRoot), subKey_(NULL)
{
    DWORD ret = ::RegCreateKeyEx(rootKey_, targetKey.c_str(), 0, NULL, REG_OPTION_VOLATILE, KEY_ALL_ACCESS, NULL,
        &subKey_, NULL);
    if (ret != ERROR_SUCCESS) {
        LOG_ERR("open hkey error: {}", ret);
    }
}

RegistryUtilsEx::~RegistryUtilsEx() noexcept
{
    if (subKey_ != NULL) {
        RegCloseKey(subKey_);
        subKey_ = NULL;
    }
}

TString RegistryUtilsEx::ToErrorString(DWORD err)
{
    TString strError;
    strError.resize(ERROR_STRING_SIZE);
    LPTSTR pStrError = const_cast<LPTSTR>(strError.c_str());
    // FormatMessage - If the function fails, the return value is zero
    if (FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, nullptr, err, 0, pStrError, ERROR_STRING_SIZE, nullptr) == 0) {
        strError = ERR_UNKNOWN;
    }
    return strError;
}

LONG RegistryUtilsEx::RegSetDword(const TString& field, const DWORD data)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    DWORD dataSize = sizeof(data);
    LOG_INFO("[RegistryUtilsEx - RegSetDword]field: {}, data: {}", field.c_str(), data);
    return ::RegSetKeyValue(subKey_, NULL, field.c_str(), REG_DWORD, (void*)&data, dataSize);
}


LONG RegistryUtilsEx::RegGetDword(const TString& field, DWORD& data)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    data = 0;
    DWORD dataSize = sizeof(data);
    return ::RegGetValue(subKey_, NULL, field.c_str(), RRF_RT_REG_DWORD, nullptr, &data, &dataSize);
}

LONG RegistryUtilsEx::RegSetString(const TString& field, const TString& data)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    DWORD dataSize = data.size() * sizeof(TCHAR);
    LOG_INFO("[RegistryUtilsEx - RegSetString]field: {}, data: {}", field.c_str(), data);
    return ::RegSetKeyValue(subKey_, NULL, field.c_str(), REG_SZ, (void*)data.c_str(), dataSize);
}

LONG RegistryUtilsEx::RegGetString(const TString& field, TString& data)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    DWORD dataSize = 0;
    LONG ret = ::RegGetValue(subKey_, NULL, field.c_str(), RRF_RT_REG_SZ, NULL, NULL, &dataSize);
    if (ret != ERROR_SUCCESS) {
        return ret;
    }

    data = STR_EMPTY;
    data.resize(dataSize / sizeof(TCHAR));
    ret = ::RegGetValue(subKey_, NULL, field.c_str(), RRF_RT_REG_SZ, NULL, &data[0], &dataSize);
    if (ret != ERROR_SUCCESS) {
        return ret;
    }

    DWORD stringLengthInWchars = dataSize / sizeof(TCHAR);
    stringLengthInWchars--; // Exclude the NUL written by the Win32 API
    data.resize(stringLengthInWchars);
    return ERROR_SUCCESS;
}

/*
 * This function has not been used. Before invoking this function, verify its correctness.
 */
LONG RegistryUtilsEx::RegGetQDword(const TString& field, ULONGLONG& data)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    data = 0;
    DWORD dataSize = sizeof(data);
    return ::RegGetValue(subKey_, NULL, field.c_str(), RRF_RT_REG_QWORD, nullptr, &data, &dataSize);
}

LONG RegistryUtilsEx::RegTraverseHKey(std::vector<std::pair<TString, DWORD>>& result)
{
    if (subKey_ == NULL) {
        LOG_ERR("error occurred, reg HKEY is null.");
        return ERROR_INVALID_HANDLE;
    }
    result.clear();
    // 给定注册表项下的值总计数和值名称最大长度
    DWORD valueCount{};
    DWORD maxValueNameLen{};
    LONG ret = ::RegQueryInfoKey(
        subKey_,
        nullptr,    // No user-defined class
        nullptr,    // No user-defined class size
        nullptr,    // Reserved
        nullptr,    // No subkey count
        nullptr,    // No subkey max length
        nullptr,    // No subkey class length
        &valueCount,
        &maxValueNameLen,
        nullptr,    // No max value length
        nullptr,    // No security descriptor
        nullptr);     // No last write time
    if (ret != ERROR_SUCCESS) {
        return ret;
    }

    // 终止符 NUL
    maxValueNameLen++;
    auto nameBuffer = std::make_unique<TCHAR[]>(maxValueNameLen);
    // Call RegEnumValue to get data of current value
    for (DWORD index = 0; index < valueCount; index++) {
        DWORD valueNameLen = maxValueNameLen;
        DWORD valueType{};
        ret = ::RegEnumValue(subKey_, index, nameBuffer.get(), &valueNameLen,
            nullptr, &valueType, nullptr, nullptr);
        if (ret != ERROR_SUCCESS) {
            continue;
        }
        result.push_back(std::make_pair(TString{ nameBuffer.get(), valueNameLen }, valueType));
    }
    return ERROR_SUCCESS;
}

}  // namespace PERF